absl-py>=0.7.0
numpy>=1.13.3
clu
jax
flax
optax
matplotlib
ml_collections
seaborn
tensorflow>=1.13.0
tensorflow_datasets